#include <d3d11.h>
#include <dxgi1_2.h>
#include <wrl/client.h>
#include <fstream>
#include <vector>
#include <chrono>
#include <thread>

using namespace Microsoft::WRL;

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "dxgi.lib")

void SaveBMP(const char* filename, uint8_t* data, int width, int height, int pitch) {
    BITMAPFILEHEADER fileHeader = {};
    BITMAPINFOHEADER infoHeader = {};

    fileHeader.bfType = 0x4D42; // 'BM'
    fileHeader.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
    fileHeader.bfSize = fileHeader.bfOffBits + pitch * height;

    infoHeader.biSize = sizeof(BITMAPINFOHEADER);
    infoHeader.biWidth = width;
    infoHeader.biHeight = -height; // top-down
    infoHeader.biPlanes = 1;
    infoHeader.biBitCount = 32;
    infoHeader.biCompression = BI_RGB;

    std::ofstream ofs(filename, std::ios::binary);
    ofs.write((char*)&fileHeader, sizeof(fileHeader));
    ofs.write((char*)&infoHeader, sizeof(infoHeader));
    ofs.write((char*)data, pitch * height);
    ofs.close();
}

int main() {
    // 1. 创建 D3D11 设备
    ComPtr<ID3D11Device> device;
    ComPtr<ID3D11DeviceContext> context;
    D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr,
        D3D11_CREATE_DEVICE_BGRA_SUPPORT, nullptr, 0,
        D3D11_SDK_VERSION, &device, nullptr, &context);

    // 2. 获取显示输出
    ComPtr<IDXGIDevice> dxgiDevice;
    device.As(&dxgiDevice);
    ComPtr<IDXGIAdapter> adapter;
    dxgiDevice->GetAdapter(&adapter);

    ComPtr<IDXGIOutput> output;
    adapter->EnumOutputs(0, &output); // 主显示器
    ComPtr<IDXGIOutput1> output1;
    output.As(&output1);

    // 3. 创建 Duplication
    ComPtr<IDXGIOutputDuplication> duplication;
    output1->DuplicateOutput(device.Get(), &duplication);

    // 获取输出信息
    DXGI_OUTPUT_DESC outputDesc;
    output->GetDesc(&outputDesc);
    UINT width = outputDesc.DesktopCoordinates.right - outputDesc.DesktopCoordinates.left;
    UINT height = outputDesc.DesktopCoordinates.bottom - outputDesc.DesktopCoordinates.top;

    // 创建 CPU 可读的 Texture
    D3D11_TEXTURE2D_DESC cpuTexDesc = {};
    cpuTexDesc.Width = width;
    cpuTexDesc.Height = height;
    cpuTexDesc.MipLevels = 1;
    cpuTexDesc.ArraySize = 1;
    cpuTexDesc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
    cpuTexDesc.SampleDesc.Count = 1;
    cpuTexDesc.Usage = D3D11_USAGE_STAGING;
    cpuTexDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    cpuTexDesc.BindFlags = 0;

    ComPtr<ID3D11Texture2D> cpuTexture;
    device->CreateTexture2D(&cpuTexDesc, nullptr, &cpuTexture);

    // 捕捉 5 帧
    for (int i = 0; i < 5; ++i) {
        DXGI_OUTDUPL_FRAME_INFO frameInfo = {};
        ComPtr<IDXGIResource> desktopResource;
        HRESULT hr = duplication->AcquireNextFrame(1000, &frameInfo, &desktopResource);
        if (FAILED(hr)) {
            printf("AcquireNextFrame failed: 0x%08X\n", hr);
            continue;
        }

        ComPtr<ID3D11Texture2D> desktopTex;
        desktopResource.As(&desktopTex);
        context->CopyResource(cpuTexture.Get(), desktopTex.Get());

        D3D11_MAPPED_SUBRESOURCE mapped;
        hr = context->Map(cpuTexture.Get(), 0, D3D11_MAP_READ, 0, &mapped);
        if (SUCCEEDED(hr)) {
            char filename[64];
            sprintf_s(filename, "frame_%d.bmp", i);
            SaveBMP(filename, (uint8_t*)mapped.pData, width, height, mapped.RowPitch);
            context->Unmap(cpuTexture.Get(), 0);
            printf("Saved %s\n", filename);
        }

        duplication->ReleaseFrame();
        std::this_thread::sleep_for(std::chrono::milliseconds(33));
    }

    return 0;
}
